package com.cn.lp.ai.factory;

import org.springframework.ai.qianfan.api.ApiUtils;
import org.springframework.ai.qianfan.api.QianFanApi;
import org.springframework.ai.qianfan.api.auth.AuthApi;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.ResponseEntity;
import org.springframework.util.Assert;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;

import java.util.function.Predicate;

public class WenxinApi extends AuthApi {

    public static final String DEFAULT_CHAT_MODEL;
    public static final String DEFAULT_EMBEDDING_MODEL;
    private static final Predicate<QianFanApi.ChatCompletionChunk> SSE_DONE_PREDICATE;
    private final RestClient restClient;
    private final WebClient webClient;

    public WenxinApi(String apiKey, String secretKey) {
        this("https://aip.baidubce.com/rpc/2.0/wenxin", apiKey, secretKey);
    }

    public WenxinApi(String baseUrl, String apiKey, String secretKey) {
        this(baseUrl, apiKey, secretKey, RestClient.builder());
    }

    public WenxinApi(String baseUrl, String apiKey, String secretKey, RestClient.Builder restClientBuilder) {
        this(baseUrl, apiKey, secretKey, restClientBuilder, RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER);
    }

    public WenxinApi(String baseUrl, String apiKey, String secretKey, RestClient.Builder restClientBuilder, ResponseErrorHandler responseErrorHandler) {
        this(baseUrl, apiKey, secretKey, restClientBuilder, WebClient.builder(), responseErrorHandler);
    }

    public WenxinApi(String baseUrl, String apiKey, String secretKey, RestClient.Builder restClientBuilder, WebClient.Builder webClientBuilder, ResponseErrorHandler responseErrorHandler) {
        super(apiKey, secretKey);
        this.restClient = restClientBuilder.baseUrl(baseUrl).defaultHeaders(ApiUtils.getJsonContentHeaders()).defaultStatusHandler(responseErrorHandler).build();
        this.webClient = webClientBuilder.baseUrl(baseUrl).defaultHeaders(ApiUtils.getJsonContentHeaders()).build();
    }

    public ResponseEntity<CustomerQianFanApiParams.TextToImageResponse> textToImage(CustomerQianFanApiParams.TextToImageRequest textToImageRequest) {
        Assert.notNull(textToImageRequest, "The request body can not be null.");
        Assert.notNull(textToImageRequest.text(), "The input can not be null.");
        Assert.isTrue(textToImageRequest.text().length() <= 190, "The text must be 190 char or less");
        return ((RestClient.RequestBodySpec)this.restClient.post().uri("/v1/basic/textToImage?access_token={token}",
                new Object[]{this.getAccessToken()})).body(textToImageRequest)
                .retrieve().toEntity(new ParameterizedTypeReference<CustomerQianFanApiParams.TextToImageResponse>() {
                });
    }

    static {
        DEFAULT_CHAT_MODEL = QianFanApi.ChatModel.ERNIE_Speed_8K.getValue();
        DEFAULT_EMBEDDING_MODEL = QianFanApi.EmbeddingModel.BGE_LARGE_ZH.getValue();
        SSE_DONE_PREDICATE = QianFanApi.ChatCompletionChunk::end;
    }

}
